1902F - Trees and XOR Queries Again - CodeForces Solution

#include<stdio.h>
#include<string.h>
#include<stdlib.h>

typedef long long LL;
#include <map>
#include <vector>
#include <queue>
#include <deque>
#include <set>
#include <stack>
#include <algorithm>
#include <array>
#include <unordered_set>
#include <unordered_map>
#include <string>
using namespace std;

#define P1 972663749
#define P2 911382323

bool rcmp(int a, int b) { return a>b; }
struct VNode {
	int v;
       	bool operator<(const VNode& b) const {
	       	return v<b.v;
       	}
};

int as[200004];
int enter[200004];
int leave[200004];
int pp[19][200004];
vector<int> nei[200004];
int gx;
int bbs[19][200004][20];
int lvl[200004];
void dfs(int r, int p, int l) {
	enter[r]=gx++;
	pp[0][r]=p;
	lvl[r]=l;
	for (auto b: nei[r]) {
		if (b==p) continue;
		dfs(b, r, l+1);
	}
	leave[r]=gx;
}

int bb[20];
void merge(int *b, int *a) {
	int k, i, v;
	for (k=0; k<20; k++) if (a[k]!=-1) {
		if (b[k]==-1) b[k]=a[k];
		else {
			v=a[k]; v^=b[k];
			for (i=k+1; v>0&&i<20; i++) if ((1<<i)&v) {
				if (b[i]==-1) { b[i]=v; break; }
				v^=b[i];
			}
		}
	}
}
int main() {
       	int n, i, a, b, x, k, ik;
	int q, na, nb;
	scanf("%d", &n);
	for (i=0; i<n; i++) scanf("%d", &as[i]);
	for (i=1; i<n; i++) {
		scanf("%d %d", &a, &b); a--; b--;
		nei[a].push_back(b);
		nei[b].push_back(a);
	}
	gx=0; dfs(0, 0, 0);
	for (i=0; i<n; i++) {
		for (k=0; k<20; k++) bbs[0][i][k]=-1;
		if (as[i]==0) continue;
		for (k=0; k<20; k++) if ((1<<k)&as[i]) break;
		bbs[0][i][k]=as[i];
	}
	for (k=1; k<19; k++) {
		for (i=0; i<n; i++) {
			pp[k][i]=pp[k-1][pp[k-1][i]];
			for (x=0; x<20; x++) bbs[k][i][x]=bbs[k-1][i][x];
			merge(bbs[k][i], bbs[k-1][pp[k-1][i]]);
		}
	}
	scanf("%d", &q);
	for (i=0; i<q; i++) {
		scanf("%d %d %d", &a, &b, &x); a--; b--;
		for (k=0; k<20; k++) bb[k]=-1;
		if (enter[b]>=enter[a]&&enter[b]<leave[a]) {}
		else {
			// raise a
			for (ik=18; ik>=0; ik--) {
				na=pp[ik][a];
				if (enter[b]>=enter[na]&&enter[b]<leave[na]) continue;
				merge(bb, bbs[ik][a]); a=na;
			}
		       	merge(bb, bbs[0][a]);
			a=pp[0][a];
		}
		merge(bb, bbs[0][a]);
		// a is lca
		// printf("lca is %d\n", a);
		if (a!=b) {
			// raise b;
			for (ik=18; ik>=0; ik--) {
				nb=pp[ik][b];
				if (lvl[nb]<=lvl[a]) continue;
				merge(bb, bbs[ik][b]);
			       	b=nb;
			}
		       	merge(bb, bbs[0][b]);
		}
		for (k=0; k<20; k++) if (bb[k]!=-1&&(x&(1<<k))) x^=bb[k];
		// for (k=0; k<20; k++) printf("%d ", bb[k]); printf("\n");
		if (x==0) printf("YES\n"); else printf("NO\n");
	}
       	return 0;
}